# ✅ 必须在最开始就替换 jieba,在任何其他导入之前
import sys
try:
    import jieba_fast
    sys.modules['jieba'] = jieba_fast
except ImportError:
    pass

# 然后才导入其他模块
from funasr import AutoModel
import numpy as np
import sounddevice as sd
import threading
import concurrent.futures


def _load_asr_model():
    """加载语音识别模型(独立线程)"""
    print("正在加载语音识别模型...")
    model = AutoModel(
        model="iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
        disable_pbar=True,
        disable_log=True,
        disable_update=True
    )
    print("语音识别模型加载完成。")
    return model


def _load_punc_model():
    """加载标点符号模型(独立线程)"""
    print("正在加载标点符号模型...")
    model = AutoModel(
        model="ct-punc",
        disable_pbar=True,
        disable_log=True,
        disable_update=True
    )
    print("标点符号模型加载完成。")
    return model


class RealtimeSpeechRecognizer:
    """
    实时语音识别类
    
    使用FunASR模型进行流式语音识别和标点添加
    """
    
    def __init__(self, sample_rate=16000, chunk_size=None, 
                 encoder_chunk_look_back=4, decoder_chunk_look_back=1):
        """
        初始化识别器
        
        Args:
            sample_rate: 采样率,默认16000Hz
            chunk_size: 分块大小 [0, 10, 5]
            encoder_chunk_look_back: 编码器回看块数
            decoder_chunk_look_back: 解码器回看块数
        """
        self.sample_rate = sample_rate
        self.chunk_size = chunk_size or [0, 10, 5]
        self.encoder_chunk_look_back = encoder_chunk_look_back
        self.decoder_chunk_look_back = decoder_chunk_look_back
        self.chunk_stride = self.chunk_size[1] * 960
        
        self._load_models()
    
    def _load_models(self):
        """并行加载语音识别和标点符号模型"""
        import logging
        
        # 设置日志级别
        if 'jieba' in sys.modules:
            sys.modules['jieba'].setLogLevel(logging.ERROR)
        
        print("开始并行加载模型...")
        
        # ✅ 使用线程池并行加载两个模型
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
            # 提交两个加载任务
            asr_future = executor.submit(_load_asr_model)
            punc_future = executor.submit(_load_punc_model)
            
            # 等待两个任务完成并获取结果
            self.model = asr_future.result()
            self.punc_model = punc_future.result()
        
        print("✅ 所有模型加载完成!")
    
    def recognize_chunk(self, audio_data, cache, is_final=False):
        """
        识别音频块
        
        Args:
            audio_data: 音频数据数组
            cache: 缓存字典,保存识别状态
            is_final: 是否为最后一块音频
            
        Returns:
            识别出的文本,无结果返回空字符串
        """
        speech_chunk = np.array(audio_data, dtype=np.float32)
        
        res = self.model.generate(
            input=speech_chunk,
            cache=cache,
            is_final=is_final,
            chunk_size=self.chunk_size,
            encoder_chunk_look_back=self.encoder_chunk_look_back,
            decoder_chunk_look_back=self.decoder_chunk_look_back
        )
        
        if res and len(res) > 0 and 'text' in res[0]:
            return res[0]['text']
        return ""
    
    def add_punctuation(self, text):
        """
        添加标点符号
        
        Args:
            text: 无标点的文本
            
        Returns:
            带标点的文本
        """
        if not text:
            return text
        
        punc_res = self.punc_model.generate(input=text)
        if punc_res and len(punc_res) > 0 and 'text' in punc_res[0]:
            return punc_res[0]['text']
        return text


class AudioRecorder:
    """
    音频录制类(单例模式)
    
    确保同一时间只有一个录音实例,避免麦克风资源冲突
    """
    
    _instance = None
    _lock = threading.Lock()
    
    def __new__(cls, *args, **kwargs):
        """单例模式实现:确保只创建一个实例"""
        if not cls._instance:
            with cls._lock:
                if not cls._instance:
                    cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(self, sample_rate=16000):
        """
        初始化录音器(只初始化一次)
        
        Args:
            sample_rate: 采样率,默认16000Hz
        """
        if hasattr(self, '_initialized'):
            return
        
        self.sample_rate = sample_rate
        self.chunk_stride = 9600  # 识别块大小
        self.audio_buffer = []  # 音频缓冲区
        self.is_recording = False  # 录音状态
        self.stream = None  # 音频流
        self.text_callback = None  # 文本回调函数
        self.cache = {}  # 识别缓存
        self.full_text = ""  # 累积的完整文本
        self.recognizer = RealtimeSpeechRecognizer()
        self._initialized = True
        self._buffer_lock = threading.Lock()
    
    def _audio_callback(self, indata, frames, time, status):
        """
        音频流回调函数(每0.1秒调用一次)
        
        Args:
            indata: 输入音频数据
            frames: 帧数
            time: 时间信息
            status: 状态信息
        """
        if not self.is_recording:
            return
        
        with self._buffer_lock:
            self.audio_buffer.extend(indata.flatten())
            
            # 当缓冲区足够大时进行识别
            if len(self.audio_buffer) >= self.chunk_stride:
                chunk = self.audio_buffer[:self.chunk_stride]
                self.audio_buffer = self.audio_buffer[self.chunk_stride:]
                
                # 实时识别音频块
                text = self.recognizer.recognize_chunk(chunk, self.cache, is_final=False)
                if text and self.text_callback:
                    self.full_text += text
                    # 推送部分识别结果
                    self.text_callback({"type": "partial", "text": text})
    
    def start_recording(self, text_callback):
        """
        开始录音
        
        Args:
            text_callback: 文本回调函数,接收识别结果
        """
        self.text_callback = text_callback
        self.is_recording = True
        self.cache = {}
        self.full_text = ""
        
        with self._buffer_lock:
            self.audio_buffer = []
        
        if self.stream is None:
            self.stream = sd.InputStream(
                samplerate=self.sample_rate,
                channels=1,
                dtype='float32',
                callback=self._audio_callback,
                blocksize=int(self.sample_rate * 0.1)
            )
            self.stream.start()
    
    def stop_recording(self):
        """
        停止录音
        
        处理剩余音频,添加标点符号,推送最终结果
        """
        self.is_recording = False
        
        def process_final():
            """后台线程处理最终结果,避免阻塞主线程"""
            remaining_audio = []
            with self._buffer_lock:
                if self.audio_buffer:
                    remaining_audio = self.audio_buffer[:]
                    self.audio_buffer = []
            
            # 处理剩余音频
            if remaining_audio:
                text = self.recognizer.recognize_chunk(
                    remaining_audio, 
                    self.cache, 
                    is_final=True
                )
                if text:
                    self.full_text += text
            
            # 清空实时显示,推送最终结果
            if self.full_text and self.text_callback:
                self.text_callback({"type": "partial", "text": ""})
                
                # 添加标点符号
                final_text = self.recognizer.add_punctuation(self.full_text)
                # 推送最终结果
                self.text_callback({"type": "final", "text": final_text})
        
        # 在独立线程中处理,不阻塞主线程
        threading.Thread(target=process_final, daemon=True).start()